#!/usr/bin/env python
import numpy as np
import itk
import scipy.ndimage as ndi
import sys
import os
import imageTools.ITKUtils.io as io
from . import utils
def anatomyExtraction(fileName, th_low, th_high, sel, suffix=".nii.gz"):
    """extract the body region from the CT image"""
    foutput = utils.getModifiedFileName(fileName,th_low,th_high)

    imgIn = io.readImage(fileName,returnITK=True,imgMode="float")
    btf = itk.BinaryThresholdImageFilter.IF3ISS3.New()
    btf.SetInput(imgIn)
    btf.SetLowerThreshold(th_low)
    btf.SetUpperThreshold(th_high)
    ccf = itk.ConnectedComponentImageFilter.ISS3IUL3.New()
    ccf.SetInput(btf.GetOutput())
    rcf = itk.RelabelComponentImageFilter.IUL3ISS3.New()
    rcf.SetInput(ccf.GetOutput())
    btf2 = itk.BinaryThresholdImageFilter.ISS3ISS3.New()
    btf2.SetInput(rcf.GetOutput())
    btf2.SetLowerThreshold(1)
    btf2.SetUpperThreshold(sel)

    img = itk.PyBuffer.ISS3.GetArrayFromImage(btf2.GetOutput())
    img = np.where(img>0,1,0).astype(np.uint8)
    newImg = np.zeros(img.shape,np.uint8)
    for i in xrange(newImg.shape[0]):
      ns = ndi.binary_fill_holes(img[i,:,:])
      newImg[i,:,:] = ns

    anatomyImg = itk.PyBuffer.ISS3.GetImageFromArray(newImg)
    io.writeImage(anatomyImg,fname=foutput,refimg=imgIn)
    return foutput

def main():
  if (len(sys.argv) < 4):
     print ("Not enough arguments")
     sys.exit(1) 

  print sys.argv[1]
  inputFile = sys.argv[1]
  th_low = int(sys.argv[2])
  th_high = int(sys.argv[3])
  sel = int(sys.argv[4])
  anatomyExtraction(inputFile, th_low, th_high, sel)

if __name__ == '__main__':
    main()
 
